package dmap

import (
    "github.com/coreos/etcd/clientv3"
    "path"
    "strings"
    "sync"
    "fmt"
    "context"
    "time"
    "github.com/coreos/etcd/mvcc/mvccpb"
)

type DistributeMap interface {
    Init() error
    Add(string) error
    Watch()
    Delete(value string) error
    Reload() error
    Query(string) (string, bool, bool)
    Close() error
}



type EtcdDistributeMap struct {
    mu         sync.RWMutex
    BaseName   string
    UID        string
    AllItems   map[string]string //value uid
    LocalItems map[string]bool
    //
    cli       *clientv3.Client
    endpoints []string
    ttl       int64
    leaseID   clientv3.LeaseID
    uname     string
    passwd    string
    cancel    context.CancelFunc
}

func NewDistributeMap(baseName, uid, uname, passwd string, endpoints []string, ttl int64) (*EtcdDistributeMap, error) {
    if uid == "" {
        uid = "watcher"
    }
    if baseName == "" || uid == "" || len(endpoints) == 0 {
        return nil, fmt.Errorf("param error")
    }
    if ttl == 0 {
        ttl = 3
    }
    return &EtcdDistributeMap{
        BaseName:  baseName,
        UID:       uid,
        endpoints: endpoints,
        ttl:       ttl,
        uname:     uname,
        passwd:    passwd,
    }, nil
}

func (m *EtcdDistributeMap) Value2Path(value string) string {
    return path.Join("/", m.BaseName, m.UID, value)
}

func (m *EtcdDistributeMap) Path2Value(p string) (uid, value string) {
    prefix := path.Join("/", m.BaseName) + "/"
    if !strings.HasPrefix(p, prefix) {
        return
    }
    s := strings.TrimPrefix(p, prefix)
    items := strings.SplitN(s, "/", 2)
    if len(items) != 2 {
        return
    }
    uid = items[0]
    value = items[1] //value without initial /
    return
}

func (m *EtcdDistributeMap) Init() error {
    cli, err := clientv3.New(clientv3.Config{
        Endpoints:   m.endpoints,
        DialTimeout: 5 * time.Second,
        Username:    m.uname,
        Password:    m.passwd,
    })
    if err != nil {
        return err
    }
    resp, err := cli.Grant(context.TODO(), m.ttl)
    if err != nil {
        return err
    }
    m.leaseID = resp.ID
    _, err = cli.KeepAlive(context.TODO(), m.leaseID)
    if err != nil {
        return err
    }
    m.mu.Lock()
    defer m.mu.Unlock()
    m.cli = cli
    m.AllItems = make(map[string]string) //value uid
    m.LocalItems = make(map[string]bool)
    return nil
}

func (m *EtcdDistributeMap) Add(value string) error {
    if m.UID == "watcher" {
        return fmt.Errorf("Ready Only Mode")
    }
    m.mu.Lock()
    defer m.mu.Unlock()
    ctx := context.TODO()
    p := m.Value2Path(value)
    v := fmt.Sprintf("%d", time.Now().Unix())
    _, err := m.cli.Put(ctx, p, v, clientv3.WithLease(m.leaseID))
    m.LocalItems[value] = true
    return err
}

func (m *EtcdDistributeMap) Watch() {
    prefix := path.Join("/", m.BaseName, m.UID) + "/"
    ctx, cancel := context.WithCancel(context.Background())
    m.cancel = cancel
    ch := m.cli.Watch(ctx, prefix, clientv3.WithPrefix())
LOOP:
    for {
        select {
        case resp := <-ch:
            if resp.Canceled {
                break LOOP
            }
            if err := resp.Err(); err != nil {
                break LOOP
            }
            for _, e := range resp.Events {
                p := string(e.Kv.Key)
                uid, value := m.Path2Value(p)
                if uid == "" || value == "" {
                    continue
                }
                m.mu.Lock()
                if e.Type == mvccpb.PUT {
                    m.AllItems[value] = uid
                } else if e.Type == mvccpb.DELETE {
                    if _, ok := m.AllItems[value]; !ok {
                        delete(m.AllItems, value)
                    }
                }
                m.mu.Unlock()
            }
        case <-ctx.Done():
            break LOOP
        }
    }
}

func (m *EtcdDistributeMap) Delete(value string) error {
    m.mu.Lock()
    defer m.mu.Unlock()
    ctx := context.TODO()
    p := m.Value2Path(value)
    _, err := m.cli.Delete(ctx, p)
    if _, ok := m.LocalItems[value]; ok {
        delete(m.LocalItems, value)
    }
    return err
}
func (m *EtcdDistributeMap) Reload() error {
    prefix := path.Join("/", m.BaseName) + "/"
    resp, err := m.cli.Get(context.TODO(), prefix, clientv3.WithPrefix())
    if err != nil {
        return err
    }
    m.mu.Lock()
    defer m.mu.Unlock()
    m.AllItems = make(map[string]string)
    m.LocalItems = make(map[string]bool)
    for _, kv := range resp.Kvs {
        p := string(kv.Key)
        uid, value := m.Path2Value(p)
        if uid == "" || value == "" {
            continue
        }
        m.AllItems[value] = uid
        if uid == m.UID {
            m.LocalItems[value] = true
        }
    }
    return nil
}

func (m *EtcdDistributeMap) Query(value string) (string, bool, bool) {
    m.mu.RLock()
    defer m.mu.RUnlock()
    uid, ok := m.AllItems[value]
    _, isLocal := m.LocalItems[value]
    return uid, ok, isLocal
}

func (m *EtcdDistributeMap) Close() error {
    if m.cancel != nil {
        m.cancel()
    }
    if m.cli != nil {
        return m.cli.Close()
    }
    return nil
}
